
from typing import AsyncGenerator, Union, Tuple
from .base import Sender
from ..schemas import Message, MessageChunk, MessageToolCall, ModelResponse
from .logger import get_logger

logger = get_logger(__name__)
class OutputAdaptor:
    def __init__(self, sender: Sender):
        self.sender=sender
        self.response=None
        self.status=None
    
    async def output(self, response: AsyncGenerator[MessageChunk, None]):
        """Internal function for streaming generation"""
        self.status="start"
        response_id=None
        provider=None
        model = None 
        role = None
        created_at=None

        _content = []
        tool_call_chunks = []
        
        logger.info(f"start to output MessageChunk")
        async for chunk in response:
            delta = chunk.choice.delta
            content_chunk = delta.content or ""
            finish_reason = chunk.choice.finish_reason

            if response_id is None:
                response_id=chunk.id
                provider = chunk.provider
                model = chunk.model
                role = chunk.role
                created_at=chunk.created_at
                
            tool_calls = {}
            if delta.tool_calls:
                for  tc in delta.tool_calls:
                    # 解析工具调用片段
                    """_summary_
                    # tc.id != "" supporting value=None is compatible with Ollama
                    >>> x=None
                    >>> if x!="": print(x)
                    ... 
                    None
                    """
                    if tc.id != "":
                        tool_call=MessageToolCall.model_validate(tc)
                        # 如果是新的工具调用，创建新对象
                        if tc.id not in tool_calls:
                            tool_calls[tc.id ] =tool_call
                            tool_call_chunks.append(tool_call)
                        else:
                            # 累积参数
                            if tc.function is not None and tc.function.arguments:
                                tool_calls[tc.id].function.arguments += tc.function.arguments
                            else: pass
                    else: pass
            # 累积完整内容
            if content_chunk:
                _content.append(content_chunk)
            # 生成统一格式的块
            await self.sender.send(chunk)
            yield chunk
            # 如果收到结束信号则退出
            if finish_reason:
                logger.info(f"{provider}/{model}: finished output MessageChunk with reason={finish_reason}")
                break

        content="".join(_content)
        thinking, assistant_text="",""
        if content != "":
            thinking, assistant_text=self.parse_content(content)
        message = Message(
                id=f"index-0",
                provider=provider,
                model=model,
                role=role,
                content=assistant_text,
                thinking=thinking,
                tool_calls=tool_call_chunks,
                created_at = created_at,
            )
        logger.info(f"{provider}/{model}: return result format response with stream")
        self.status="end"
        self.response= ModelResponse(
                id=response_id,
                provider=message.provider,
                created_at=message.created_at,
                message=message,
                metrics=None )

    def parse_content(self, content: str)-> Tuple[str,str]:
        start=content.find(f"<think>")
        end=content.find(f"</think>")
        thinking=""
        assistant_text=""
        if start !=-1 and end !=-1:
            thinking=content[start+7:end].strip()
            assistant_text=content[end+8:].strip()
        else:
            assistant_text=content.strip() 
        return thinking, assistant_text